import pandas as pd
import numpy as np
from datetime import datetime

def load_data():
    """加载用户和购买数据"""
    try:
        customers = pd.read_csv('customer_data.csv')
        purchases = pd.read_csv('purchase_history.csv')
        
        # 转换日期列
        purchases['purchase_date'] = pd.to_datetime(purchases['purchase_date'])
        
        return customers, purchases
    except FileNotFoundError:
        print("找不到数据文件，请先运行data_generator.py生成数据")
        return None, None

def create_features(customers_df, purchases_df):
    """创建用于预测的特征"""
    if customers_df is None or purchases_df is None:
        return None
    
    # 计算每个用户的总购买次数
    purchase_counts = purchases_df.groupby('customer_id')['purchase_date'].count().reset_index()
    purchase_counts.columns = ['customer_id', 'total_purchases']
    
    # 计算每个用户的平均购买金额
    avg_spend = purchases_df.groupby('customer_id')['amount'].mean().reset_index()
    avg_spend.columns = ['customer_id', 'avg_purchase_amount']
    
    # 计算每个用户的总消费金额
    total_spend = purchases_df.groupby('customer_id')['amount'].sum().reset_index()
    total_spend.columns = ['customer_id', 'total_spend']
    
    # 计算每个用户的首次购买日期和最近购买日期
    first_purchase = purchases_df.groupby('customer_id')['purchase_date'].min().reset_index()
    first_purchase.columns = ['customer_id', 'first_purchase_date']
    
    last_purchase = purchases_df.groupby('customer_id')['purchase_date'].max().reset_index()
    last_purchase.columns = ['customer_id', 'last_purchase_date']
    
    # 计算每个用户的复购次数
    repurchase_counts = purchases_df[purchases_df['is_repurchase'] == True].groupby('customer_id')['purchase_date'].count().reset_index()
    repurchase_counts.columns = ['customer_id', 'repurchase_count']
    
    # 合并所有特征
    features = customers_df.merge(purchase_counts, on='customer_id', how='left')
    features = features.merge(avg_spend, on='customer_id', how='left')
    features = features.merge(total_spend, on='customer_id', how='left')
    features = features.merge(first_purchase, on='customer_id', how='left')
    features = features.merge(last_purchase, on='customer_id', how='left')
    features = features.merge(repurchase_counts, on='customer_id', how='left')
    
    # 填充可能的NaN值
    features['total_purchases'] = features['total_purchases'].fillna(0)
    features['avg_purchase_amount'] = features['avg_purchase_amount'].fillna(0)
    features['total_spend'] = features['total_spend'].fillna(0)
    features['repurchase_count'] = features['repurchase_count'].fillna(0)
    
    # 创建目标变量：是否为复购用户（至少有一次复购）
    features['is_repeat_buyer'] = features['repurchase_count'] > 0
    
    # 计算其他特征
    today = datetime.now()
    features['days_since_first_purchase'] = (today - features['first_purchase_date']).dt.days
    features['days_since_last_purchase'] = (today - features['last_purchase_date']).dt.days
    features['purchase_frequency'] = features['total_purchases'] / (features['days_since_first_purchase'] + 1) * 30  # 每月购买频率
    
    # 计算RFM特征
    features['recency'] = features['days_since_last_purchase']
    features['frequency'] = features['total_purchases']
    features['monetary'] = features['total_spend']
    
    # 处理分类特征
    features = pd.get_dummies(features, columns=['gender', 'product_category'])
    
    return features

def save_features(features_df):
    """保存特征数据到CSV文件"""
    if features_df is not None:
        features_df.to_csv('features.csv', index=False)
        print("特征工程完成，数据已保存到features.csv")

if __name__ == "__main__":
    customers, purchases = load_data()
    features = create_features(customers, purchases)
    save_features(features)    